from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.backends import default_backend


def generate_rsa_keys():
    """生成RSA密钥对"""
    private_key = rsa.generate_private_key(
        public_exponent=65537, key_size=2048, backend=default_backend()
    )
    public_key = private_key.public_key()
    return private_key, public_key


def rsa_encrypt(public_key, plaintext):
    """使用公钥加密数据"""
    # 确保明文是字节类型
    if isinstance(plaintext, str):
        plaintext = plaintext.encode("utf-8")

    ciphertext = public_key.encrypt(
        plaintext,
        padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None,
        ),
    )
    return ciphertext


def rsa_decrypt(private_key, ciphertext):
    """使用私钥解密数据"""
    plaintext = private_key.decrypt(
        ciphertext,
        padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None,
        ),
    )
    # 将解密后的字节转换为字符串
    return plaintext.decode("utf-8")


if __name__ == "__main__":
    # 生成RSA密钥对
    private_key, public_key = generate_rsa_keys()
    print("已生成RSA密钥对")

    # 要加密的原始数据
    original_data = "这是一个RSA加解密的示例，包含中文和特殊字符！@#$%^&*()"
    print(f"\n原始数据: {original_data}")

    # 使用公钥加密
    encrypted_data = rsa_encrypt(public_key, original_data)
    print(f"\n加密后的数据: {encrypted_data.hex()}")  # 以十六进制显示

    # 使用私钥解密
    decrypted_data = rsa_decrypt(private_key, encrypted_data)
    print(f"\n解密后的数据: {decrypted_data}")

    # 验证解密结果
    if original_data == decrypted_data:
        print("\n验证成功：解密后的数据与原始数据一致")
    else:
        print("\n验证失败：解密后的数据与原始数据不一致")
